from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_socketio import SocketIO, emit, join_room, leave_room
import pymysql
import random
import smtplib
from email.mime.text import MIMEText
from email.header import Header
import hashlib
import time
import jwt
import datetime

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

# JWT 和 SocketIO 配置
SECRET_KEY = "your-secret-key"  # 请替换为安全的密钥
app.config['SECRET_KEY'] = SECRET_KEY
socketio = SocketIO(app, cors_allowed_origins="*")

# 固定专家的用户名
EXPERT_USERNAME = "123456"

# 数据库连接配置
def get_db_connection():
    return pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='dachuang',
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

# 生成6位验证码
def generate_verification_code():
    return ''.join([str(random.randint(0, 9)) for _ in range(6)])

# 发送QQ邮箱验证码
def send_verification_email(email, code):
    sender = '952304392@qq.com'
    password = 'lwlbpkrftgzxbbfb'

    msg = MIMEText(f'您的验证码是：{code}，有效期20分钟', 'plain', 'utf-8')
    msg['From'] = Header(sender)
    msg['To'] = Header(email)
    msg['Subject'] = Header('验证码')

    try:
        smtp_obj = smtplib.SMTP_SSL('smtp.qq.com', 465)
        smtp_obj.login(sender, password)
        smtp_obj.sendmail(sender, [email], msg.as_string())
        smtp_obj.quit()
        return True
    except Exception as e:
        print(f"邮件发送失败: {e}")
        return False

# 密码哈希
def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

# 生成 JWT Token
def generate_token(user_id, username):
    payload = {
        'user_id': user_id,
        'username': username,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=24)
    }
    return jwt.encode(payload, SECRET_KEY, algorithm='HS256')

# 验证 JWT Token
def verify_token(token):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=['HS256'])
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

# 请求注册验证码
@app.route('/register/code', methods=['POST'])
def register_code():
    data = request.get_json()
    username = data.get('username')
    email = data.get('email')

    if not all([username, email]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT * FROM users WHERE username = %s OR email = %s', (username, email))
    if c.fetchone():
        conn.close()
        return jsonify({'message': '用户名或邮箱已存在'}), 400

    code = generate_verification_code()

    if send_verification_email(email, code):
        c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                  'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                  (email, code, int(time.time()), code, int(time.time())))
        conn.commit()
        conn.close()
        return jsonify({'message': '验证码已发送至您的邮箱'})
    else:
        conn.close()
        return jsonify({'message': '验证码发送失败'}), 500

# 注册接口
@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '验证码无效'}), 400

    stored_code, timestamp = result['code'], result['timestamp']

    if int(time.time()) - timestamp > 1200:
        conn.close()
        return jsonify({'message': '验证码已过期'}), 400

    if stored_code != code:
        conn.close()
        return jsonify({'message': '验证码错误'}), 400

    hashed_password = hash_password(password)
    c.execute('INSERT INTO users (username, password, email, verified) VALUES (%s, %s, %s, %s)',
              (username, hashed_password, email, 1))
    conn.commit()
    c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
    conn.commit()
    conn.close()

    return jsonify({'message': '注册成功'})

# 请求登录验证码
@app.route('/login/code', methods=['POST'])
def login_code():
    data = request.get_json()
    email = data.get('email')

    if not email:
        return jsonify({'message': '缺少邮箱参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT * FROM users WHERE email = %s', (email,))
    if not c.fetchone():
        conn.close()
        return jsonify({'message': '邮箱未注册'}), 400

    code = generate_verification_code()

    if send_verification_email(email, code):
        c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                  'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                  (email, code, int(time.time()), code, int(time.time())))
        conn.commit()
        conn.close()
        return jsonify({'message': '验证码已发送至您的邮箱'})
    else:
        conn.close()
        return jsonify({'message': '验证码发送失败'}), 500

# 登录接口
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT id, password, email FROM users WHERE username = %s', (username,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '用户不存在'}), 400

    user_id, stored_password, stored_email = result['id'], result['password'], result['email']

    if stored_email != email:
        conn.close()
        return jsonify({'message': '邮箱不匹配'}), 400

    if hash_password(password) != stored_password:
        conn.close()
        return jsonify({'message': '密码错误'}), 400

    c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '验证码无效'}), 400

    stored_code, timestamp = result['code'], result['timestamp']

    if int(time.time()) - timestamp > 1200:
        conn.close()
        return jsonify({'message': '验证码已过期'}), 400

    if stored_code != code:
        conn.close()
        return jsonify({'message': '验证码错误'}), 400

    c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
    conn.commit()
    conn.close()

    token = generate_token(user_id, username)
    return jsonify({
        'message': '登录成功',
        'token': token,
        'user_id': user_id,
        'username': username
    })

# 获取所有用户（专家端）
@app.route('/users', methods=['GET'])
def get_users():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:]
    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    if payload['username'] != EXPERT_USERNAME:
        return jsonify({'message': '仅专家可访问'}), 403

    conn = get_db_connection()
    c = conn.cursor()
    c.execute('SELECT id, username FROM users WHERE username != %s', (EXPERT_USERNAME,))
    users = c.fetchall()
    conn.close()

    return jsonify([{'id': u['id'], 'username': u['username']} for u in users])

# 获取聊天历史
@app.route('/messages/<int:receiver_id>', methods=['GET'])
def get_messages(receiver_id):
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:]
    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    user_id = payload['user_id']
    username = payload['username']

    conn = get_db_connection()
    c = conn.cursor()

    if username != EXPERT_USERNAME:
        c.execute('SELECT id FROM users WHERE username = %s', (EXPERT_USERNAME,))
        expert = c.fetchone()
        if not expert or receiver_id != expert['id']:
            conn.close()
            return jsonify({'message': '用户只能与专家聊天'}), 403

    c.execute('''
        SELECT sender_id, receiver_id, content, timestamp
        FROM messages
        WHERE (sender_id = %s AND receiver_id = %s) OR (sender_id = %s AND receiver_id = %s)
        ORDER BY timestamp ASC
    ''', (user_id, receiver_id, receiver_id, user_id))
    messages = c.fetchall()
    conn.close()

    return jsonify([{
        'sender_id': m['sender_id'],
        'receiver_id': m['receiver_id'],
        'content': m['content'],
        'timestamp': m['timestamp'].isoformat()
    } for m in messages])

# WebSocket 事件
@socketio.on('connect')
def handle_connect():
    token = request.args.get('token')
    if not token:
        return False
    payload = verify_token(token)
    if not payload:
        return False
    emit('status', {'message': '连接成功'})

@socketio.on('join')
def on_join(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    user_id = payload['user_id']
    room = f'user_{user_id}'
    join_room(room)
    emit('status', {'message': f'用户 {user_id} 已加入房间 {room}'}, room=room)

@socketio.on('leave')
def on_leave(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    user_id = payload['user_id']
    room = f'user_{user_id}'
    leave_room(room)
    emit('status', {'message': f'用户 {user_id} 已离开房间 {room}'}, room=room)

@socketio.on('message')
def handle_message(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    sender_id = payload['user_id']
    username = payload['username']
    receiver_id = data.get('receiver_id')
    content = data.get('content')

    if not all([receiver_id, content]):
        emit('error', {'message': '缺少必要参数'})
        return

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT id FROM users WHERE id = %s', (receiver_id,))
    if not c.fetchone():
        conn.close()
        emit('error', {'message': '接收者不存在'})
        return

    if username != EXPERT_USERNAME:
        c.execute('SELECT id FROM users WHERE username = %s', (EXPERT_USERNAME,))
        expert = c.fetchone()
        if not expert or receiver_id != expert['id']:
            conn.close()
            emit('error', {'message': '用户只能与专家聊天'})
            return

    c.execute('INSERT INTO messages (sender_id, receiver_id, content) VALUES (%s, %s, %s)',
              (sender_id, receiver_id, content))
    conn.commit()

    c.execute('SELECT timestamp FROM messages WHERE id = LAST_INSERT_ID()')
    timestamp = c.fetchone()['timestamp']
    conn.close()

    receiver_room = f'user_{receiver_id}'
    sender_room = f'user_{sender_id}'
    message_data = {
        'sender_id': sender_id,
        'receiver_id': receiver_id,
        'content': content,
        'timestamp': timestamp.isoformat()
    }
    emit('message', message_data, room=receiver_room)
    emit('message', message_data, room=sender_room)

if __name__ == '__main__':
    socketio.run(app, debug=True, host='0.0.0.0', port=5000, allow_unsafe_werkzeug=True)